# -*- coding: utf-8 -*-

from sympy import S, Abs, solveset, Interval
from mathsolver.functions.base.base import BaseFunction, new_latex
from mathsolver.functions.base.objects import BaseFunc
from mathsolver.functions.hanshu.dingyiyu import QiuDingYiYu
from mathsolver.functions.hanshu.helper import check_func, elim_re
from mathsolver.functions.sympy_utils import default_symbol


class QuJueDuiZhi(BaseFunction):
    """
    递归去绝对值
    1. 不含绝对值
    2. 最底层的绝对值
    3. 绝对值嵌套
    4. 其他含绝对值的情况
    """

    @property
    def get_piecefunc_expr(self):
        return [[biaodashi, qujian] for qujian, biaodashi in self.output]

    # 从给定区间(伪)随机取数,从而决定去绝对值后的正负号
    @staticmethod
    def get_value_from_interval(interval):
        ll = interval.args[0]
        r = interval.args[1]
        if ll.is_finite and r.is_finite:
            return (ll + r) / 2
        elif ll.is_finite and not r.is_finite:
            return ll + 1
        elif not ll.is_finite and r.is_finite:
            return r - 1
        else:
            return 1

    def elim_abs(self, f, interval=S.Reals):
        target_symbol = default_symbol(f)
        if "Abs" not in str(f):  # 1. 不含绝对值
            return interval, [(interval, f, f)]

        elif isinstance(f, Abs) and "Abs" not in str(f.args[0]):  # 2. 最底层的绝对值
            expr_in_abs = f.args[0]
            zeros = solveset(expr_in_abs, target_symbol, domain=S.Reals)  # 求出零点

            # 补充：求 expr_in_Abs 的定义域
            dingyiyu = QiuDingYiYu().solver(BaseFunc({'name': 'f', 'var': target_symbol, 'type': '', 'expression': expr_in_abs}))
            interval = interval.intersect(dingyiyu)

            sub_intervals = zeros.complement(interval)  # 用零点把定义域分成若干子区间
            interval_exp_pairs = []
            # 考虑到只有一个子区间的情况
            ll = sub_intervals.args[0]
            if ll.is_number:
                test_value = self.get_value_from_interval(sub_intervals)
                if expr_in_abs.subs(target_symbol, test_value) > 0:
                    interval_exp_pairs.append((sub_intervals, f, expr_in_abs))
                else:
                    interval_exp_pairs.append((sub_intervals, f, -expr_in_abs))
            else:  # 多个子区间
                for sub in sub_intervals.args:
                    test_value = self.get_value_from_interval(sub)
                    if expr_in_abs.subs(target_symbol, test_value) > 0:
                        interval_exp_pairs.append((sub, f, expr_in_abs))
                    else:
                        interval_exp_pairs.append((sub, f, -expr_in_abs))
            return sub_intervals, interval_exp_pairs

        elif isinstance(f, Abs) and "Abs" in str(f.args[0]):  # 3. 绝对值嵌套
            intervals = interval
            expr_without_abs = []
            sub_intervals, pair = self.elim_abs(f.args[0], interval)
            intervals = intervals.intersect(sub_intervals)
            # 考虑到只有一个子区间的情况
            ll = intervals.args[0]
            if ll.is_number:
                pai = pair[0]
                to_be_elim = f.subs(pai[1], pai[2])
                sub_sub_interval, sub_pair = self.elim_abs(to_be_elim, sub_intervals)
                intervals = sub_intervals.complement(intervals).union(sub_sub_interval)
                for v in sub_pair:
                    expr_without_abs.append((v[0], f, v[2]))  # 把要替换的项v[1]改成f
            else:  # 多个子区间
                for sub_interval, pai in zip(sub_intervals.args, pair):
                    to_be_elim = f.subs(pai[1], pai[2])
                    sub_sub_interval, sub_pair = self.elim_abs(to_be_elim, sub_interval)
                    intervals = sub_interval.complement(intervals).union(sub_sub_interval)
                    for v in sub_pair:
                        expr_without_abs.append((v[0], f, v[2]))  # 把要替换的项v[1]改成f
            return intervals, expr_without_abs

        elif "Abs" in str(f):  # 4. 其他含绝对值的情况
            pairs = []
            intervals = interval
            for arg in f.args:
                if "Abs" in str(arg):
                    sub_intervals, pair = self.elim_abs(arg, interval)
                    pairs.append(pair)
                    intervals = intervals.intersect(sub_intervals)
            expr_without_abs = []
            # 考虑到只有一个子区间的情况
            ll = intervals.args[0]
            if ll.is_number:
                f1 = f
                for pair in pairs:
                    for sub_int in pair:
                        if intervals.is_subset(sub_int[0]):  # 是 sub_int[0] 的子集
                            f1 = f1.subs(sub_int[1], sub_int[2])
                            break
                expr_without_abs.append((intervals, f, f1))
            else:  # 多个子区间
                for sub_interval in intervals.args:
                    f1 = f
                    for pair in pairs:
                        for sub_int in pair:
                            if sub_interval.is_subset(sub_int[0]):  # 是 sub_int[0] 的子集
                                f1 = f1.subs(sub_int[1], sub_int[2])
                                break

                    expr_without_abs.append((sub_interval, f, f1))

            return intervals, expr_without_abs

    def solver(self, *args):
        """
        去绝对值
        :param args: 第一个参数为BaseEq, 第二个参数为取值区间, 可选参数, 默认为R
        :return: self.output
        """
        func = check_func(args[0])
        # 补充：sympify('Abs(x*E**x)') => exp(re(x))*Abs(x)
        f = elim_re(func.expression)

        if len(args) == 2:
            interval = args[1]
        else:
            interval = S.Reals
        _, result = self.elim_abs(f, interval)

        self.steps.append(["", "原式 = $ %s $" % new_latex(f)])
        self.steps.append(["在区间 %s" % interval, "去绝对值后得到分段函数"])
        for i, v in enumerate(result):
            subinterval = v[0]
            if i > 0:
                subinterval = Interval(v[0].args[0], v[0].args[1], False, v[0].args[3])
            self.steps.append(["$%s" % new_latex(subinterval), "\\quad\t%s$" % new_latex(v[2])])
            self.output.append([subinterval, v[2]])

        self.output = self.get_piecefunc_expr
        self.label.add("去绝对值")
        return self


if __name__ == '__main__':
    pass
